package org.JWCB.judge.config;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.io.FileUtil;
import cn.hutool.core.util.ArrayUtil;
import cn.hutool.core.util.StrUtil;
import com.github.dockerjava.api.DockerClient;
import com.github.dockerjava.api.command.*;
import com.github.dockerjava.api.model.*;
import org.JWCB.common.core.constants.JudgeConstants;
import org.JWCB.judge.callback.DockerStartResultCallback;
import org.springframework.beans.factory.annotation.Value;

import java.io.File;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;



public class DockerSandBoxPool {

    private DockerClient dockerClient;

    private String sandboxImage;

    private String volumeDir;

    private Long memoryLimit;

    private Long memorySwapLimit;

    private Long cpuLimit;

    private int poolSize;

    private String containerNamePrefix;

    private BlockingQueue<String> containerIdQueue;

    private Map<String,String> containerIdNameMap;


    private String mountDir;

    public DockerSandBoxPool(String containerNamePrefix, int poolSize, Long cpuLimit, Long memorySwapLimit, Long memoryLimit, String volumeDir, String sandboxImage, DockerClient dockerClient,String mountDir) {
        this.containerNamePrefix = containerNamePrefix;
        this.poolSize = poolSize;
        this.cpuLimit = cpuLimit;
        this.memorySwapLimit = memorySwapLimit;
        this.memoryLimit = memoryLimit;
        this.volumeDir = volumeDir;
        this.sandboxImage = sandboxImage;
        this.dockerClient = dockerClient;
        this.containerIdQueue = new ArrayBlockingQueue<>(poolSize);
        this.containerIdNameMap = new HashMap<>();
        this.mountDir = mountDir;
    }


    /**
     * 初始化容器
     */
    public void initDockerPool() {
        for (int i = 0; i < poolSize; i++) {
            String containerName = containerNamePrefix + "-" + i;
            String containerId = createContainer(containerName);
        }
    }


    /**
     * 获取容器
     */
    public String getContainer() {
        try {
            return containerIdQueue.take();
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
    }


    public void returnContainer(String containerId) {
        containerIdQueue.add(containerId);
    }
    /**
     * 创建容器
     */
    private String createContainer(String containerName) {
        /**
         * 先判定容器是否已经存在
         */
        List<Container> containerList = dockerClient.listContainersCmd().withShowAll(true).exec();
        if(!CollectionUtil.isEmpty(containerList)) {
            String names = JudgeConstants.JAVA_CONTAINER_PREFIX + containerName;
            for (Container container : containerList) {
                String[] containerNames = container.getNames();
                if(containerNames != null && containerNames.length > 0 && names.equals(containerNames[0])) {
                    if("created".equals(container.getState()) || "exited".equals(container.getState())) {
                        dockerClient.startContainerCmd(container.getId()).exec();
                    }
                    containerIdQueue.add(container.getId());
                    containerIdNameMap.put(container.getId(),containerName);
                    return container.getId();
                }
            }
        }
        // 拉取java环境镜像
        pullJavaEnvImage();
        HostConfig hostConfig = getHostConfig(containerName);
        CreateContainerCmd createContainerCmd = dockerClient
                .createContainerCmd(JudgeConstants.JAVA_ENV_IMAGE)
                .withName(containerName);
        CreateContainerResponse createConfigResponse = createContainerCmd
                .withHostConfig(hostConfig)
                .withAttachStderr(true)
                .withAttachStderr(true)
                .withTty(true)
                .exec();
        String containerId = createConfigResponse.getId();

        dockerClient.startContainerCmd(containerId).exec();
        containerIdQueue.add(containerId);
        containerIdNameMap.put(containerId,containerName);
        return containerId;
    }



    /**
     * 拉取java环境镜像
     */
    private void pullJavaEnvImage() {
        ListImagesCmd listImagesCmd = dockerClient.listImagesCmd();
        List<Image> imageList = listImagesCmd.exec();
        for(Image image : imageList) {
            String[] repoTags = image.getRepoTags();
            if(repoTags != null && repoTags.length != 0 && sandboxImage.equals(repoTags[0])) {
                return;
            }
        }
        PullImageCmd pullImageCmd = dockerClient.pullImageCmd(sandboxImage);
        try {
            pullImageCmd.exec(new PullImageResultCallback().awaitCompletion());
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
    }



    /**
     * 限制资源
     */
    private HostConfig getHostConfig(String containerName) {
        HostConfig hostConfig = new HostConfig();
        String userCodeDir = createContainerDir(containerName);
        hostConfig.setBinds(new Bind(userCodeDir,new Volume(volumeDir)));
        hostConfig.withMemory(memoryLimit);
        hostConfig.withMemorySwap(memorySwapLimit);
        hostConfig.withCpuCount(cpuLimit);
        hostConfig.withNetworkMode("none");
        hostConfig.withReadonlyRootfs(true);
        return hostConfig;
    }


    /**
     * 创建挂载目录
     */
    private String createContainerDir(String containerName) {
        String codeDir = System.getProperty("user.dir")  + "/" + JudgeConstants.CODE_DIR_POOL;
        if (!FileUtil.exist(codeDir)) {
            FileUtil.mkdir(codeDir);
        }
        return codeDir + "/" + containerName;
    }


    /**
     * 获取容器的挂载目录
     */

    public String getCodeDir(String containerId) {
        String containerName = containerIdNameMap.get(containerId);
        return System.getProperty("user.dir")  + "/" + JudgeConstants.CODE_DIR_POOL + containerName;
    }


}
